# -*- coding: utf-8 -*- 
# @Time : 2022/4/3 20:55 
# @Author : zzuxyj 
# @File : 10-nn-sequentialTest.py

"""
PyTorch 中有一些基础概念在构建网络的时候很重要，比如 nn.Module, nn.ModuleList, nn.Sequential，这些类我们称之为容器 (containers)，因为我们可以添加模块 (module) 到它们之中。
"""
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = Sequential(
            # (in_size - kernel_size + stride + 2padding ) / stride = out_size
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            # stride 默认 kernel_size大小  ; ceil_model 模式就是会把不足square_size的边给保留下来，单独另算，或者也可以理解为在原来的数据上补充了值为-NAN的边。而floor模式则是直接把不足square_size的边给舍弃了。
            MaxPool2d(kernel_size=2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            # 改变形状  二维张量，形状通常为[batch_size, size] [64,16]
            Flatten(),  # 64 个 一维 且长度 为 4*4 = 16
            # 全连接层 4*4*64= 1024  输出out_channel即也是神经元数量
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        output = self.model(input)
        return output


# 训练
def train(writer , input):
    model = Model()
    writer.add_graph(model , input)
    # 输出模型看看
    print(model)
    output = model(input)
    print(output.shape)


if __name__ == '__main__':

    testSet = CIFAR10("../dataset/CIFAR10", download=True, train=False, transform=ToTensor())
    dataloader = DataLoader(testSet, batch_size=64)
    # tensorboard
    writer = SummaryWriter("logs10")
    step = 0
    for data in dataloader:
        imgs, target = data
        output = train(writer , imgs)
        break

"""
Model(
  (model): Sequential(
    (0): Conv2d(3, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (4): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Flatten(start_dim=1, end_dim=-1)
    (7): Linear(in_features=1024, out_features=64, bias=True)
    (8): Linear(in_features=64, out_features=10, bias=True)
  )
)
torch.Size([64, 10])

Process finished with exit code 0

"""